/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stddef.h>

#include "mbedtls/aes.h"
#include "mbedtls/cipher.h"
#include "mbedtls/gcm.h"
#include "mbedtls/hkdf.h"
#include "mbedtls/md.h"
#include "mbedtls/md_internal.h"
#include "mbedtls/sha256.h"
#include "securec.h"
#include "ola_def.h"
#include "ola_sdk.h"
#include "ola_base_utils.h"
#include "ola_security.h"
#include "ola_error.h"
#include "ola_pal_os.h"
#include "ola_pal_store.h"
#include "ola_log.h"
#include "ola_random.h"
#include "ola_data.h"

#define OLA_DATA_SALT_LEN 16
#define OLA_DATA_KEY_LEN 32

enum {
    OLA_DATA_SEC_LEVEL_LOW,
    OLA_DATA_SEC_LEVEL_HIGH
};

typedef struct {
    unsigned char *buf;
    const char *key;
    unsigned int len;
    unsigned int size;
    unsigned int level;
} OlaSecData;

static OlaSecData g_olaSecData[] = {
    { NULL, "ticket", OLA_DEV_TICKET_LEN + 1, 0, OLA_DATA_SEC_LEVEL_LOW },
    { NULL, "devid", OLA_DEV_REGISTER_DEVID_LEN + 1, 0, OLA_DATA_SEC_LEVEL_LOW },
    { NULL, "secret", OLA_DEV_REGISTER_SECRET_LEN + 1, 0, OLA_DATA_SEC_LEVEL_HIGH },
    { NULL, "couldurl", OLA_DEV_CLOUD_URL_LEN + 1, 0, OLA_DATA_SEC_LEVEL_LOW },
    { NULL, "registrycode", OLA_DEV_REGISTRY_CODE_LEN + 1, 0, OLA_DATA_SEC_LEVEL_LOW },
    { NULL, "wififlag", OLA_DEV_WIFI_INFO_FLAG_LEN + 1, 0, OLA_DATA_SEC_LEVEL_LOW },
    { NULL, "psk", OLA_DEV_WIFI_INFO_FLAG_LEN + 1, 0, OLA_DATA_SEC_LEVEL_HIGH }
};

static bool g_olaDataInitFlag = false;

static int OlaDataGetKeyMaterial(unsigned char *out, unsigned int outLen)
{
    OlaData key = {NULL, 0};

    if (OlaGetDeviceData(OLA_DATA_TYPE_STORAGE_KEY_MATERIAL, &key) != OLA_OK) {
        OlaLogError("get key material failed");
        return OLA_API_ERR_GET_DATA;
    }

    if ((key.len == 0) || (key.data == NULL)) {
        OlaLogError("invalid key material data");
        return OLA_API_ERR_GET_DATA;
    }

    if (memcpy_s(out, outLen, key.data, key.len) != EOK) {
        OlaLogError("memcpy_s failed");
        OlaDataFree(&key);
        return OLA_SECUREC_ERR_MEMCOPY;
    }

    OlaDataFree(&key);
    return OLA_OK;
}

static int OlaDataGetKey(unsigned char *key, unsigned int keyLen, const unsigned char *salt, unsigned int saltLen)
{
    unsigned char tempKey[OLA_DATA_KEY_LEN] = {0};
    int ret;

    if (keyLen < OLA_DATA_KEY_LEN) {
        return OLA_DATA_ERR_INVALID_LEN;
    }

    ret = OlaDataGetKeyMaterial(tempKey, sizeof(tempKey));
    if (ret != OLA_OK) {
        return ret;
    }

    const mbedtls_md_info_t *md = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
    ret = mbedtls_md_hmac(md, tempKey, sizeof(tempKey), salt, saltLen, key);
    if (ret != 0) {
        OlaLogError("get key failed");
        (void)memset_s(tempKey, sizeof(tempKey), 0, sizeof(tempKey));
        return OLA_DATA_ERR_GET_KEY;
    }

    (void)memset_s(tempKey, sizeof(tempKey), 0, sizeof(tempKey));
    return OLA_OK;
}

static int OlaDataEncode(OlaData *input, OlaData *key, OlaData *output)
{
    if ((input->data == NULL) || (input->len == 0)) {
        OlaLogWarning("encode, in data invalid");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    /* 密文组包形式|ver(1byte)|iv(12bytes)|data|tag(16bytes)|salt(16bytes)(预留空间外部填充)| */
    output->len = input->len + OLA_TYPICAL_GCM_VERSION_LEN +
        OLA_TYPICAL_GCM_IV_LEN + OLA_TYPICAL_GCM_TAG_LEN + OLA_DATA_SALT_LEN;
    output->data = OlaMalloc(output->len);
    if (output->data == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }
    (void)memset_s(output->data, output->len, 0, output->len);
    output->data[0] = OLA_TYPICAL_GCM_VERSION;

    int ret;
    do {
        ret = OlaRand(output->data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN);
        if (ret != OLA_OK) {
            OlaLogError("gen iv failed");
            break;
        }

        AesGcmParam inputParam = {
            .key = {key->data, key->len},
            .iv = {output->data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN},
            .add = {NULL, 0},
            .tag = {NULL, 0},
            .inData = {input->data, input->len},
        };

        OlaBuffer encData = {output->data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN, 0, input->len};
        OlaBuffer tag = {output->data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN + input->len,
            0, OLA_TYPICAL_GCM_TAG_LEN};
        ret = OlaAesGcmEncript(&inputParam, &encData, &tag);
        if (ret != OLA_OK) {
            OlaLogWarning("aes-256-gcm encript error");
            break;
        }
        return OLA_OK;
    } while (0);

    OlaDataFree(output);
    return ret;
}

static int OlaDataDecode(const OlaData *input, const OlaData *key, OlaData *output)
{
    if (input->len <= (OLA_TYPICAL_GCM_IV_LEN + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_TAG_LEN)) {
        OlaLogWarning("decode, in data invalid");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    /* 密文组包形式|ver(1byte)|iv(12bytes)|data|tag(16bytes)| */
    AesGcmParam inputParam = {
        .key = {key->data, key->len},
        .iv = {input->data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN},
        .add = {NULL, 0},
        .tag = {
            input->data + input->len - OLA_TYPICAL_GCM_TAG_LEN,
            OLA_TYPICAL_GCM_TAG_LEN
        },
        .inData = {
            input->data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN,
            input->len - OLA_TYPICAL_GCM_VERSION_LEN - OLA_TYPICAL_GCM_IV_LEN - OLA_TYPICAL_GCM_TAG_LEN
        },
    };

    output->data = OlaMalloc(inputParam.inData.len);
    if (output->data == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }

    OlaBuffer outBuffer = {output->data, 0, inputParam.inData.len};
    int ret = OlaAesGcmDecript(&inputParam, &outBuffer);
    if (ret != OLA_OK) {
        OlaDataFree(output);
        outBuffer.buf = NULL;
    } else {
        output->len = outBuffer.len;
    }

    return ret;
}

static int OlaPreSetData(OlaDataType type, const OlaData *in)
{
    if ((in == NULL) || (in->data == NULL)) {
        OlaLogError("invalid input");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (type >= OLA_ARRAY_SIZE(g_olaSecData)) {
        OlaLogError("invalid type");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if ((in->len == 0) || (in->len >= g_olaSecData[type].len)) {
        OlaLogError("invalid input len");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (g_olaSecData[type].level == OLA_DATA_SEC_LEVEL_HIGH) {
        return OLA_OK;
    }

    if (memcpy_s(g_olaSecData[type].buf, g_olaSecData[type].len, in->data, in->len) != EOK) {
        OlaLogError("memcpy_s failed");
        return OLA_SECUREC_ERR_MEMCOPY;
    }
    g_olaSecData[type].size = in->len;
    return OLA_OK;
}

int OlaSetData(OlaDataType type, OlaData *in)
{
    unsigned char salt[OLA_DATA_SALT_LEN] = {0};
    unsigned char key[OLA_DATA_KEY_LEN] = {0};
    OlaData out = { NULL, 0 };
    int ret;

    if (!g_olaDataInitFlag) {
        return OLA_DATA_ERR_NOT_INIT;
    }

    OlaData keyData = { key, OLA_DATA_KEY_LEN };
    ret = OlaPreSetData(type, in);
    if (ret != OLA_OK) {
        goto OLA_SET_DATA_EXIT;
    }

    ret = OlaRand(salt, OLA_DATA_SALT_LEN);
    if (ret != OLA_OK) {
        OlaLogError("get rand failed");
        goto OLA_SET_DATA_EXIT;
    }

    ret = OlaDataGetKey(key, OLA_DATA_KEY_LEN, salt, OLA_DATA_SALT_LEN);
    if (ret != OLA_OK) {
        OlaLogError("get key failed");
        goto OLA_SET_DATA_EXIT;
    }

    ret = OlaDataEncode(in, &keyData, &out);
    if (ret != OLA_OK) {
        goto OLA_SET_DATA_EXIT;
    }

    if (memcpy_s(out.data + out.len - OLA_DATA_SALT_LEN, OLA_DATA_SALT_LEN, salt, OLA_DATA_SALT_LEN) != EOK) {
        OlaLogError("copy salt failed");
        ret = OLA_SECUREC_ERR_MEMCOPY;
        goto OLA_SET_DATA_EXIT;
    }

    ret = OlaStoreSet(g_olaSecData[type].key, out.data, out.len);
    if (ret != 0) {
        OlaLogError("write data failed");
        ret = OLA_DATA_ERR_RAW_WRITE;
        goto OLA_SET_DATA_EXIT;
    }
    OlaLogInfo("set [%s] len [%u] sucess", g_olaSecData[type].key, out.len);
OLA_SET_DATA_EXIT:
    (void)memset_s(key, OLA_DATA_KEY_LEN, 0, OLA_DATA_KEY_LEN);
    OlaDataFree(&out);
    return ret;
}

static int OlaPreLoadData(unsigned char index, OlaData *in, unsigned char *salt)
{
    unsigned int len = OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN +
        OLA_TYPICAL_GCM_TAG_LEN + OLA_DATA_SALT_LEN;

    in->len = len + g_olaSecData[index].len;
    in->data = (unsigned char *)OlaMalloc(in->len);
    if (in->data == NULL) {
        OlaLogError("malloc failed");
        return OLA_OS_ERR_MALLOC;
    }
    (void)memset_s(in->data, in->len, 0, in->len);

    if (OlaStoreGet(g_olaSecData[index].key, in->data, &(in->len)) != OLA_OK) {
        OlaLogInfo("read data failed");
        return OLA_DATA_ERR_RAW_READ;
    }
    OlaLogInfo("get [%s] len [%u] sucess", g_olaSecData[index].key, in->len);
    if (in->len <= len) {
        OlaLogError("read len failed");
        return OLA_DATA_ERR_INVALID_LEN;
    }

    if (memcpy_s(salt, OLA_DATA_SALT_LEN, in->data + in->len - OLA_DATA_SALT_LEN, OLA_DATA_SALT_LEN) != EOK) {
        OlaLogError("memcpy_s failed");
        return OLA_SECUREC_ERR_MEMCOPY;
    }
    in->len = in->len - OLA_DATA_SALT_LEN;

    return OLA_OK;
}

static int OlaLoadData(unsigned char index, OlaData *output)
{
    unsigned char salt[OLA_DATA_SALT_LEN] = {0};
    unsigned char key[OLA_DATA_KEY_LEN] = {0};
    OlaData keyData = { key, OLA_DATA_KEY_LEN };
    OlaData in = { NULL, 0 };
    OlaData out = { NULL, 0 };
    int ret;

    ret = OlaPreLoadData(index, &in, salt);
    if (ret != OLA_OK) {
        goto OLA_DATA_LOAD_EXIT;
    }

    ret = OlaDataGetKey(key, OLA_DATA_KEY_LEN, salt, OLA_DATA_SALT_LEN);
    if (ret != OLA_OK) {
        OlaLogError("get key failed");
        goto OLA_DATA_LOAD_EXIT;
    }

    ret = OlaDataDecode(&in, &keyData, &out);
    if (ret != OLA_OK) {
        goto OLA_DATA_LOAD_EXIT;
    }

    if (g_olaSecData[index].level == OLA_DATA_SEC_LEVEL_HIGH) {
        ret = memcpy_s(output->data, output->len, out.data, out.len);
    } else {
        ret = memcpy_s(g_olaSecData[index].buf, g_olaSecData[index].len, out.data, out.len);
    }

    if (ret != EOK) {
        OlaLogError("memcpy_s failed");
        ret = OLA_SECUREC_ERR_MEMCOPY;
        goto OLA_DATA_LOAD_EXIT;
    }

    if (g_olaSecData[index].level == OLA_DATA_SEC_LEVEL_HIGH) {
        output->len = out.len;
    } else {
        g_olaSecData[index].size = out.len;
    }

OLA_DATA_LOAD_EXIT:
    (void)memset_s(key, OLA_DATA_KEY_LEN, 0, OLA_DATA_KEY_LEN);
    OlaDataFree(&in);
    OlaDataFree(&out);
    return ret;
}

int OlaGetData(OlaDataType type, OlaData *out)
{
    if (out == NULL) {
        OlaLogError("invalid para");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (!g_olaDataInitFlag) {
        return OLA_DATA_ERR_NOT_INIT;
    }

    if (type >= OLA_ARRAY_SIZE(g_olaSecData)) {
        OlaLogError("invalid type");
        return OLA_COM_ERR_NOT_SUPPORT;
    }

    if (g_olaSecData[type].level == OLA_DATA_SEC_LEVEL_LOW) {
        out->data = g_olaSecData[type].buf;
        out->len = g_olaSecData[type].size;
        return OLA_OK;
    }

    if ((out->data == NULL) || (out->len == 0)) {
        OlaLogError("invalid para");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    return OlaLoadData(type, out);
}

bool OlaIsDevAuth()
{
    if (g_olaSecData[OLA_DATA_TYPE_TICKET].buf == NULL) {
        return false;
    }

    if (g_olaSecData[OLA_DATA_TYPE_TICKET].size == 0) {
        return false;
    }

    return true;
}

bool OlaIsRegister(void)
{
    if (g_olaSecData[OLA_DATA_TYPE_DEV_ID].buf == NULL) {
        return false;
    }

    if (g_olaSecData[OLA_DATA_TYPE_DEV_ID].size == 0) {
        return false;
    }

    unsigned char secret[OLA_DEV_REGISTER_SECRET_LEN + 1] = {0};
    OlaData secData = { secret, OLA_DEV_REGISTER_SECRET_LEN };
    if (OlaLoadData(OLA_DATA_TYPE_SECRET, &secData) != OLA_OK) {
        (void)memset_s(secret, sizeof(secret), 0, sizeof(secret));
        return false;
    }

    if (secData.len == 0) {
        (void)memset_s(secret, sizeof(secret), 0, sizeof(secret));
        return false;
    }

    (void)memset_s(secret, sizeof(secret), 0, sizeof(secret));
    return true;
}

int OlaRestoreFactorySettings(void)
{
    unsigned char i;

    for (i = 0; i < OLA_ARRAY_SIZE(g_olaSecData); i++) {
        (void)OlaStoreDel(g_olaSecData[i].key);
        OlaFree(g_olaSecData[i].buf);
        g_olaSecData[i].buf = NULL;
        g_olaSecData[i].size = 0;
    }

    return OLA_OK;
}

int OlaDataInit()
{
    unsigned char i;

    if (g_olaDataInitFlag) {
        return OLA_DATA_ERR_ALREADY_INIT;
    }

    for (i = 0; i < OLA_ARRAY_SIZE(g_olaSecData); i++) {
        if (g_olaSecData[i].level == OLA_DATA_SEC_LEVEL_HIGH) {
            continue;
        }

        g_olaSecData[i].buf = (unsigned char *)OlaMalloc(g_olaSecData[i].len);
        if (g_olaSecData[i].buf == NULL) {
            OlaLogError("malloc failed");
            continue;
        }
        (void)memset_s(g_olaSecData[i].buf, g_olaSecData[i].len, 0, g_olaSecData[i].len);

        (void)OlaLoadData(i, NULL);
    }

    g_olaDataInitFlag = true;
    return OLA_OK;
}

int OlaDataDeinit()
{
    unsigned char i;

    if (!g_olaDataInitFlag) {
        return OLA_DATA_ERR_NOT_INIT;
    }

    for (i = 0; i < OLA_ARRAY_SIZE(g_olaSecData); i++) {
        if (g_olaSecData[i].buf != NULL) {
            OlaFree(g_olaSecData[i].buf);
            g_olaSecData[i].buf = NULL;
            g_olaSecData[i].size = 0;
        }
    }

    g_olaDataInitFlag = false;
    return OLA_OK;
}
